import random
import time

import sympy
import math
import matplotlib.pyplot as pl
from mpl_toolkits.mplot3d import Axes3D as ax3
import numpy as np
from numpy.linalg import LinAlgError


'''
6. 对于随机生成的(P, q, r, x 0 )，编程实现“梯度下降+精确直线搜索”求解上述优化问题
7. 对于n = 2，在随机给定初始(P, q, r, x 0 )的情况下（建议尝试多组不同的初始值(P, q, r, x 0 )）， 结合上述程序，在平面上绘制点列 { x k } k ≥ 0 的图像
8. 对于较大的n值（例如n = 100甚至更大）， 在随机给定初始(P, q, r, x 0 )的情况下（建 议尝试多组不同的初始值(P, q, r, x 0 )），结合上述程序， 讨论点列(x 0 , x 1 , ...)的收敛速
度（即迭代次数）以及程序的运行速度（与直接计算x ∗ = −P −1 q相比较）
'''


def getP(dim):
    while 1:
        A = np.random.rand(dim, dim)
        B = np.dot(A, A.transpose())
        C = B+B.T  # makesure symmetric
        try:
            # test whether C is definite
            np.linalg.cholesky(C)  # if there is no error, C is definite
            return C
        except ValueError:
            pass
        except LinAlgError:
            pass


# 最速下降法
def SD(x0, P, q, r, N, E):
    f = lambda x: 0.5 * (np.dot(np.dot(x.T, P), x)) + np.dot(q.T, x) + r  # 定义原函数
    f_d = lambda x: np.dot(P, x) + q  # 导数
    x_draw = []
    y_draw = []

    X = x0
    Y = []  # 存储y的值
    Y_d = []  # 存储y的导数值
    xx = sympy.symarray('xx', (2, 1))
    n = 1
    ee = f_d(x0)

    e = (ee[0] ** 2 + ee[1] ** 2) ** 0.5  # 评价迭代精度
    Y.append(f(x0)[0, 0])
    Y_d.append(e)
    λ = sympy.Symbol('λ', real=True)  # 步长因子
    print('第%d次迭代：e=%d' % (n, e))
    while n < N and e > E:
        start = time.time()
        n = n + 1
        yy = f(x0 - λ * f_d(x0))
        yy_d = sympy.diff(yy[0, 0], λ, 1)  # 求导,使用sympy.diff函数，传入2个参数：函数表达式和变量名
        a0 = sympy.solve(yy_d)  # 使用sympy.solve函数解方程
        x0 = x0 - a0 * f_d(x0)  # 更新迭代
        X = np.c_[X, x0]
        Y.append(f(x0)[0, 0])
        ee = f_d(x0)
        e = math.pow(math.pow(ee[0, 0], 2) + math.pow(ee[1, 0], 2), 0.5)
        Y_d.append(e)
        end = time.time()
        print('第%d次迭代：e=%s,lr=%s, run_time=%s' % (n, e, str(a0), str(end-start)))
        print("梯度：")
        print(ee)

    return X, Y, Y_d


# 6.编程实现"梯度下降+精确直线搜索"
if __name__ == '__main__':
    n = 100
    P = getP(n)
    q = np.random.randn(n, 1)
    r = random.random()
    print("P :")
    print(P)
    print("qt :")
    print(q.T)
    print("r :")
    print(r)
    f = lambda x: 0.5 * (np.dot(np.dot(x.T, P), x)) + np.dot(q.T, x) + r

    f_d = lambda x: np.dot(P, x) + q

    x0 = np.random.randn(n, 1)
    print("x0 :")
    print(x0)
    print("f_d(x0)")
    print(f_d(x0))
    N = 1000
    E = 10 ** (-3)
    X, Y, Y_d = SD(x0, P, q, r, N, E)
    # print("X:")
    # print(X)
    # print("Y:")
    # print(Y)
    # print("Y_d:")
    # print(Y_d)

# 作业7 对于n==2时，画图
    if n == 2:
        # 画图
        X = np.array(X)
        Y = np.array(Y)
        Y_d = np.array(Y_d, dtype=object)
        figure1 = pl.figure('trend')
        n = len(Y)
        x = np.arange(1, n + 1)
        pl.subplot(2, 1, 1)
        pl.semilogy(x, Y, 'r*')
        pl.xlabel('n')
        pl.ylabel('f(x)')
        pl.subplot(2, 1, 2)
        pl.semilogy(x, Y_d, 'g*')
        pl.xlabel('n')
        pl.ylabel("|f'(x)|")
        figure2 = pl.figure('behave')
        x = np.arange(-110, 110, 1)
        y = x
        [xx, yy] = np.meshgrid(x, y)
        zz = np.zeros(xx.shape)
        n = xx.shape[0]
        for i in range(n):
            for j in range(n):
                xxx = np.array([xx[i, j], yy[i, j]])
                zz[i, j] = f(xxx.T)
        ax = ax3(figure2)
        ax.contour3D(xx, yy, zz)
        ax.plot3D(X[0, :], X[1, :], Y, 'ro--')
        pl.show()

        # ---------------------{xk的散点图}--------------------
        pic_x1 = []
        pic_x2 = []
        for i in range(len(Y)):
            pic_x1.append(X[0][i])
            pic_x2.append(X[1][i])
        pl.scatter(pic_x1, pic_x2)
        pl.xlabel("x[0]")
        pl.ylabel("x[1]")
        print(len(pic_x1))
        pl.show()
